package io.longpolling.support;

import cn.hutool.core.thread.NamedThreadFactory;
import com.alibaba.fastjson.JSON;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;

import javax.servlet.AsyncContext;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Collection;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

/**
 * 异步响应任务.
 */
@Slf4j
public class AsyncResponseTask {

    /**
     * servlet-异步请求上下文.
     * <p>持有当前请求的request和response对象.</p>
     */
    private final AsyncContext asyncContext;
    private final RequestParam requestParam;

    /**
     * 响应超时返回内容.
     */
    private final TimeoutResult timeoutResult;

    private volatile boolean complete = false; //当前任务是否已完成响应
    private final ScheduledFuture<?> timeoutTaskFuture; //超时任务引用


    public AsyncResponseTask(AsyncContext asyncContext, RequestParam requestParam, TimeoutResult timeoutResult) {

        this.asyncContext = asyncContext;
        this.timeoutResult = timeoutResult;
        this.requestParam = requestParam;

        offerToTaskPool(requestParam.getResource(), this);
        this.timeoutTaskFuture = registerTimeoutTask();
    }

    public static void buildAndStart(HttpServletRequest request, HttpServletResponse response, RequestParam requestParam, TimeoutResult timeoutResult) {
        AsyncContext asyncContext = request.startAsync(request, response); //开启异步请求上下文
        new AsyncResponseTask(asyncContext, requestParam, timeoutResult);
    }

    /**
     * 响应json内容.
     * <p>header[ContentType:application/json;charset=UTF-8]</p>
     *
     * @param result 响应数据内容
     * @param code   response-header状态码
     */
    public void returnJsonResult(Object result, HttpStatus code) {
        if (!isComplete()) {
            try {
                final HttpServletResponse response = (HttpServletResponse) asyncContext.getResponse();
                response.setStatus(code.value());
                response.setContentType(MediaType.APPLICATION_JSON_UTF8_VALUE);
                response.getWriter().println(JSON.toJSONString(result));
            } catch (Exception e) {
                log.error("Async response exception! result:{}", result, e);
            } finally {
                complete();
            }
        }
    }

    public boolean isComplete() {
        return complete;
    }

    private void complete() {
        complete = true;
        try {
            asyncContext.complete(); //释放请求上下文
        } catch (Exception e) {
            //don't care
        }

        try {
            timeoutTaskFuture.cancel(true); //取消超时任务.
        } catch (Exception e) {
            //don't care
        }

        removeFromTaskPool();

    }

    private ScheduledFuture<?> registerTimeoutTask() {
        return TIMEOUT_TASK_SCHEDULER.schedule(() -> {
            returnJsonResult(timeoutResult.getResult(), timeoutResult.getStatus());
            log.debug("超时响应完成！timeoutResult:{}", timeoutResult);
        }, timeoutResult.getTimeoutSecond(), TimeUnit.SECONDS);
    }

    private void offerToTaskPool(String resource, AsyncResponseTask asyncResponseTask) {
        RESOURCE_ASYNC_TASK_POOL.put(resource, asyncResponseTask);
    }

    private void removeFromTaskPool() {
        RESOURCE_ASYNC_TASK_POOL.remove(requestParam.getResource(), this);
    }

    public RequestParam getRequestParam() {
        return requestParam;
    }

    public static Collection<AsyncResponseTask> getAsyncTask(String resource) {
        return RESOURCE_ASYNC_TASK_POOL.get(resource);
    }

    /**
     * 响应超时任务执行器.
     * <p>单线程.</p>
     */
    private static final ScheduledExecutorService TIMEOUT_TASK_SCHEDULER = Executors.newSingleThreadScheduledExecutor(new NamedThreadFactory("async-resp-task-", false));

    /**
     * 请求资源-异步响应任务池.
     * <p>Multimap：维护资源与响应任务一对多关系.</p>
     */
    private static final Multimap<String, AsyncResponseTask> RESOURCE_ASYNC_TASK_POOL = Multimaps.synchronizedMultimap(HashMultimap.create());

}
